import sys
import os

# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import json
import os
import os.path as osp
import argparse
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
import multiprocessing
from tictoc import TicToc
import shutil

def read_camera_info(json_path):
    with open(json_path, 'r') as f:
        calibs = json.load(f)
    return calibs


def trans_img_new_intrin(out_intrin, new_size, old_intrin, img, filename):
    H = np.dot(out_intrin, np.linalg.inv(old_intrin))
    img = Image.fromarray(img[:, :, ::-1])

    # img.show()

    w, h = img.size
    H_inv = np.linalg.inv(H)
    H_inv = H_inv / H_inv[2, 2]  # 归一化
    # PIL expects a 3x3 matrix for perspective transform
    coefficients = H_inv.flatten()[:8]  # 取前8个元素，用于透视变换
    # 进行仿射变换
    if isinstance(new_size, list):
        w = new_size[1]
        h = new_size[0]
    new_img = img.transform(
        (w, h),
        Image.PERSPECTIVE,
        coefficients,
        Image.BICUBIC
    )
    # new_img = self.remove_black_borders(new_img)
    # self.show_diff_intrin_imgs(img, new_img)
    # img.show()
    # new_img.show()

    new_img.save(filename)
    return new_img


def main_worker(dataset_path, args):
    out_intrin = [[1.2726e+03, 0.0000e+00, 8.2662e+02],
                  [0.0000e+00, 1.2726e+03, 4.7975e+02],
                  [0.0000e+00, 0.0000e+00, 1.0000e+00]]
    new_size = [900, 1600]
    json_path = osp.join(dataset_path, 'calib', 'calib.json')
    imgs_path = osp.join(dataset_path, 'camera')
    sensors_info = read_camera_info(json_path)
    select_cameras = ['camera75', 'camera77', 'camera80', 'camera81']

    for select_camera in select_cameras:
        try:
            # 相对软链接
            os.chdir(os.path.join(dataset_path, 'camera'))
            os.symlink(select_camera, select_camera + '_normal')
            # os.symlink(osp.join(imgs_path, select_camera), osp.join(imgs_path, select_camera + '_normal'))
            print(
                f"符号链接已创建: {osp.join(imgs_path, select_camera + '_normal')} -> {osp.join(imgs_path, select_camera + '_normal')}")
        except OSError as e:
            print(f"创建符号链接时出错: {e}")

        sensors_info[select_camera]['normal'] = dict()
        sensors_info[select_camera]['normal']['imgh'] = sensors_info[select_camera]['ori']['imgh']
        sensors_info[select_camera]['normal']['imgw'] = sensors_info[select_camera]['ori']['imgw']
        sensors_info[select_camera]['normal']['K'] = sensors_info[select_camera]['ori']['K']
    # if os.path.exists(os.path.join(dataset_path, 'camera_epe')):
    #     shutil.rmtree(os.path.join(dataset_path, 'camera_epe'))
    # if os.path.exists(os.path.join(dataset_path, 'camera_extra')):
    #     shutil.rmtree(os.path.join(dataset_path, 'camera_extra'))

    with open(json_path, 'w') as f:
        json.dump(sensors_info, f, ensure_ascii=False, indent=2)

    print(f'done')


def main(args):
    # 统计耗时
    cost = TicToc("相机归一化")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/media/adt/T7/ZWH/docker/files/data/carla",
                        help='your data root for kitti')
    args = parser.parse_args()

    main(args)
